package com.liuhengjia.langchain4japply.api;

import com.liuhengjia.langchain4japply.config.assistant.Assistant;
import com.liuhengjia.langchain4japply.config.assistant.StructuredOutputAssistant;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.request.ResponseFormat;
import dev.langchain4j.model.chat.request.ResponseFormatType;
import dev.langchain4j.model.chat.request.json.JsonObjectSchema;
import dev.langchain4j.model.chat.request.json.JsonSchema;
import dev.langchain4j.model.chat.response.ChatResponse;
import lombok.RequiredArgsConstructor;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

/**
 * 结构化输出 API
 */
@RestController
@RequestMapping("/api/structured-output")
@RequiredArgsConstructor
public class StructuredOutputAPI {
    private final ChatLanguageModel chatLanguageModel;
    private final Assistant structuredOutputAssistant;

    /**
     * 低级接口
     *
     * @param message 用户输入
     * @return 模型响应
     */
    @GetMapping("/low/chat")
    public String lowChat(@RequestParam(name = "message") String message) {
        ResponseFormat responseFormat = ResponseFormat.builder()
                .type(ResponseFormatType.JSON)
                .jsonSchema(JsonSchema.builder()
                        .rootElement(JsonObjectSchema.builder()
                                .addIntegerProperty("age", "The person's age")
                                .addIntegerProperty("weight", "The person's weight")
                                .required("age", "weight")
                                .build())
                        .build())
                .build();
        ChatResponse chatResponse = chatLanguageModel.chat(ChatRequest.builder()
                .messages(List.of(UserMessage.from(message)))
                .parameters(ChatRequestParameters.builder()
                        .responseFormat(responseFormat)
                        .build())
                .build());
        return chatResponse.aiMessage().text();
    }

    /**
     * 高级接口
     *
     * @param message 用户输入
     * @return 模型响应
     */
    @GetMapping("/high/chat")
    public String highChat(@RequestParam(name = "message") String message) {
        return ((StructuredOutputAssistant) structuredOutputAssistant).extractPerson(message).toString();
    }
}
